package com.cyx.exercise.ollama;

import java.io.IOException;
import java.util.Scanner;
import java.util.concurrent.CompletableFuture;

/**
 * 改进的Ollama示例 - 专门处理qwen3等思考型模型
 */
public class ImprovedOllamaExample {
    
    private static final String OLLAMA_BASE_URL = "http://localhost:11434";
    private static final String DEFAULT_MODEL = "qwen3:4b";
    
    public static void main(String[] args) {
        ImprovedOllamaExample example = new ImprovedOllamaExample();
        
        System.out.println("=== 改进的Ollama Java客户端示例 ===");
        System.out.println("专门优化处理qwen3等思考型模型");
        System.out.println("模型: " + DEFAULT_MODEL);
        System.out.println("服务地址: " + OLLAMA_BASE_URL);
        System.out.println();
        
        try {
            // 1. 测试连接
            example.testConnection();
            
            // 2. 测试简单问答
            example.testSimpleQA();
            
            // 3. 测试聊天功能
            example.testChat();
            
            // 4. 交互式聊天
            example.interactiveChat();
            
        } catch (Exception e) {
            System.err.println("示例执行出错: " + e.getMessage());
            e.printStackTrace();
        }
    }
    
    /**
     * 测试连接
     */
    private void testConnection() {
        System.out.println("1. 测试连接...");
        try {
            OllamaClient client = new OllamaClient(OLLAMA_BASE_URL);
            String models = client.getModels();
            System.out.println("✓ 连接成功！");
            
            // 检查是否包含我们要使用的模型
            if (models.contains(DEFAULT_MODEL)) {
                System.out.println("✓ 找到目标模型: " + DEFAULT_MODEL);
            } else {
                System.out.println("⚠ 未找到目标模型: " + DEFAULT_MODEL);
                System.out.println("可用模型列表: " + models.substring(0, Math.min(200, models.length())) + "...");
            }
        } catch (Exception e) {
            System.err.println("✗ 连接失败: " + e.getMessage());
            throw new RuntimeException("无法连接到Ollama服务", e);
        }
        System.out.println();
    }
    
    /**
     * 测试简单问答
     */
    private void testSimpleQA() {
        System.out.println("2. 测试简单问答...");
        try {
            OllamaClient client = new OllamaClient(OLLAMA_BASE_URL);
            
            // 测试数学问题
            System.out.println("问题: 1+1等于几？请直接回答。");
            OllamaRequest request = new OllamaRequest(DEFAULT_MODEL, "1+1等于几？请直接回答。");
            
            long startTime = System.currentTimeMillis();
            OllamaResponse response = client.generate(request);
            long endTime = System.currentTimeMillis();
            
            // 处理响应
            OllamaResponseProcessor.ProcessedResponse processed = 
                OllamaResponseProcessor.processResponse(response.getResponse());
            
            System.out.println("原始响应长度: " + response.getResponse().length() + " 字符");
            System.out.println("耗时: " + (endTime - startTime) + "ms");
            System.out.println("完成状态: " + response.isDone());
            
            if (processed.hasThinkingProcess()) {
                System.out.println("✓ 检测到思考过程");
                System.out.println("思考过程长度: " + processed.getThinkingProcess().length() + " 字符");
                System.out.println("实际答案: " + processed.getAnswer());
            } else {
                System.out.println("直接答案: " + processed.getDisplayText());
            }
            
        } catch (Exception e) {
            System.err.println("✗ 简单问答测试失败: " + e.getMessage());
        }
        System.out.println();
    }
    
    /**
     * 测试聊天功能
     */
    private void testChat() {
        System.out.println("3. 测试聊天功能...");
        try {
            OllamaClient client = new OllamaClient(OLLAMA_BASE_URL);
            
            OllamaChatRequest chatRequest = new OllamaChatRequest(DEFAULT_MODEL);
            chatRequest.addSystemMessage("你是一个有用的助手，请简洁地回答问题。");
            chatRequest.addUserMessage("请用一句话解释什么是Java编程语言？");
            
            System.out.println("用户: 请用一句话解释什么是Java编程语言？");
            
            long startTime = System.currentTimeMillis();
            OllamaChatResponse chatResponse = client.chat(chatRequest);
            long endTime = System.currentTimeMillis();
            
            // 处理聊天响应
            OllamaResponseProcessor.ProcessedResponse processed = 
                OllamaResponseProcessor.processResponse(chatResponse.getMessage().getContent());
            
            System.out.println("耗时: " + (endTime - startTime) + "ms");
            System.out.println("角色: " + chatResponse.getMessage().getRole());
            System.out.println("完成状态: " + chatResponse.isDone());
            
            if (processed.hasThinkingProcess()) {
                System.out.println("助手(思考后): " + processed.getAnswer());
                System.out.println("(包含 " + processed.getThinkingProcess().length() + " 字符的思考过程)");
            } else {
                System.out.println("助手: " + processed.getDisplayText());
            }
            
        } catch (Exception e) {
            System.err.println("✗ 聊天功能测试失败: " + e.getMessage());
        }
        System.out.println();
    }
    
    /**
     * 交互式聊天
     */
    private void interactiveChat() {
        System.out.println("4. 交互式聊天 (输入 'quit' 退出, 'debug' 显示详细信息)...");
        
        try {
            OllamaClient client = new OllamaClient(OLLAMA_BASE_URL);
            Scanner scanner = new Scanner(System.in);
            OllamaChatRequest chatRequest = new OllamaChatRequest(DEFAULT_MODEL);
            
            // 设置系统提示
            chatRequest.addSystemMessage("你是一个有用的AI助手，请用中文简洁地回答问题。");
            
            boolean debugMode = false;
            
            while (true) {
                System.out.print("用户: ");
                String userInput = scanner.nextLine().trim();
                
                if ("quit".equalsIgnoreCase(userInput)) {
                    System.out.println("聊天结束。");
                    break;
                }
                
                if ("debug".equalsIgnoreCase(userInput)) {
                    debugMode = !debugMode;
                    System.out.println("调试模式: " + (debugMode ? "开启" : "关闭"));
                    continue;
                }
                
                if (userInput.isEmpty()) {
                    continue;
                }
                
                // 添加用户消息
                chatRequest.addUserMessage(userInput);
                
                try {
                    long startTime = System.currentTimeMillis();
                    
                    // 异步发送请求以提供更好的用户体验
                    CompletableFuture<OllamaChatResponse> future = client.generateAsync(
                        new OllamaRequest(DEFAULT_MODEL, userInput)
                    ).thenApply(response -> {
                        // 将OllamaResponse转换为OllamaChatResponse
                        OllamaChatMessage message = new OllamaChatMessage("assistant", response.getResponse());
                        return new OllamaChatResponse(message, response.isDone());
                    });
                    
                    System.out.print("助手正在思考");
                    
                    // 显示等待动画
                    while (!future.isDone()) {
                        System.out.print(".");
                        Thread.sleep(500);
                    }
                    System.out.println();
                    
                    OllamaChatResponse response = future.get();
                    long endTime = System.currentTimeMillis();
                    
                    // 处理响应
                    OllamaResponseProcessor.ProcessedResponse processed = 
                        OllamaResponseProcessor.processResponse(response.getMessage().getContent());
                    
                    if (debugMode) {
                        System.out.println("--- 调试信息 ---");
                        System.out.println("耗时: " + (endTime - startTime) + "ms");
                        System.out.println("原始响应长度: " + response.getMessage().getContent().length() + " 字符");
                        System.out.println("有思考过程: " + processed.hasThinkingProcess());
                        if (processed.hasThinkingProcess()) {
                            System.out.println("思考过程长度: " + processed.getThinkingProcess().length() + " 字符");
                        }
                        System.out.println("--- 调试信息结束 ---");
                    }
                    
                    // 显示助手回复
                    String assistantResponse = processed.getDisplayText();
                    System.out.println("助手: " + assistantResponse);
                    
                    // 将助手的回复添加到对话历史中
                    chatRequest.addAssistantMessage(assistantResponse);
                    
                } catch (Exception e) {
                    System.err.println("聊天请求失败: " + e.getMessage());
                    if (debugMode) {
                        e.printStackTrace();
                    }
                }
                
                System.out.println();
            }
            
        } catch (Exception e) {
            System.err.println("交互式聊天失败: " + e.getMessage());
        }
    }
}
